#/usr/bin/perl

use strict;
use warnings;

use DBI;

# ---------------------------------------------------------------------------------------
# Are we past the Rosetta stage?
# Test for whether it says we're done in the DB
# Rosetta could be finished, but this will return false if it isn't "DONE" in the DB
# ---------------------------------------------------------------------------------------
sub is_rosetta_done {

	print STDERR "--- check_status.pm: is_rosetta_done\n";

	my $pdb = $_[0];
	my $chain1 = $_[1];
	my $chain2 = $_[2];

	warn "ERROR: PDB code undefined." unless defined $pdb;
	warn "ERROR: Chain 1 undefined."  unless defined $chain1;
	warn "ERROR: Chain 2 undefined."  unless defined $chain2;

	my $dbh = DBI->connect(
		"dbi:SQLite:dbname=workflow_tracker.db",
		"",
		"",
		{ RaiseError => 1 },
	) or die $DBI::errstr;

	# Grab Rosetta job status
	my $sql = "SELECT pdb_partner.rosetta_status ";
	$sql   .= "FROM pdb, pdb_partner ";
	$sql   .= "WHERE pdb.pdb_id = '$pdb' ";
	$sql   .= "AND pdb_partner.chain_1 = '$chain1' ";
	$sql   .= "AND pdb_partner.chain_2 = '$chain2' ";
	$sql   .= "AND pdb.id = pdb_partner.pdb_id;";
	
	my $sth = $dbh->prepare($sql);
	$sth->execute();

	my $ros_status = $sth->fetch();

	my $test_result = 0;
	if (defined (@$ros_status) && join ("", @$ros_status) eq "DONE") {
		$test_result = 1;
	}
	
	return $test_result;

	$sth->finish();
	$dbh->disconnect();

}

# ---------------------------------------------------------------------------------------
# Is Rosetta running?
# Search for the job ID (stored in the DB) using qstat
# ---------------------------------------------------------------------------------------
sub is_rosetta_running {

	print STDERR "--- check_status.pm: is_rosetta_running\n";

	my $pdb = $_[0];
	my $chain1 = $_[1];
	my $chain2 = $_[2];

	warn "ERROR: PDB code undefined." unless defined $pdb;
	warn "ERROR: Chain 1 undefined."  unless defined $chain1;
	warn "ERROR: Chain 2 undefined."  unless defined $chain2;

	my $dbh = DBI->connect(
		"dbi:SQLite:dbname=workflow_tracker.db",
		"",
		"",
		{ RaiseError => 1 },
	) or die $DBI::errstr;

	# Grab Rosetta job ID
	my $sql = "SELECT pdb_partner.rosetta_job_id ";
	$sql   .= "FROM pdb, pdb_partner ";
	$sql   .= "WHERE pdb.pdb_id = '$pdb' ";
	$sql   .= "AND pdb_partner.chain_1 = '$chain1' ";
	$sql   .= "AND pdb_partner.chain_2 = '$chain2' ";
	$sql   .= "AND pdb.id = pdb_partner.pdb_id;";
	
	my $sth = $dbh->prepare($sql);
	$sth->execute();

	my $ros_job_id = $sth->fetch();

	my $job_status;

	# Only do if we have a job ID
	if (@$ros_job_id[0]) {
		my $ros_job_id_str = join ("", @$ros_job_id);
	
		my $qstat_cmd = "qstat $ros_job_id_str | tail -n1 | awk '{ print \$5 }'";
		$job_status = `$qstat_cmd`;
		chomp $job_status;
	}
	
	my $test_result = 0;
	if (defined ($job_status) && $job_status eq "R") {
		$test_result = 1;
	}
	
	return $test_result;

	$sth->finish();
	$dbh->disconnect();

}

# ---------------------------------------------------------------------------------------
# Is Rosetta actually finished (successfully)?
# Check last line of ala_out file
# ---------------------------------------------------------------------------------------
sub is_rosetta_finished {

	print STDERR "--- check_status.pm: is_rosetta_finished\n";

	my $pdb = $_[0];
	my $chain1 = $_[1];
	my $chain2 = $_[2];

	warn "ERROR: PDB code undefined." unless defined $pdb;
	warn "ERROR: Chain 1 undefined."  unless defined $chain1;
	warn "ERROR: Chain 2 undefined."  unless defined $chain2;

	my $ala_out_file = 'rosetta_output/' . $pdb . '_';
	$ala_out_file .= $chain1 . '_' . $chain2 . '_ala_out.txt';

	my $test_result = 0;

	# Check if ala out file exists, otherwise just return 0
	if (-e $ala_out_file) {

		my $grep_cmd = "tail -n1 $ala_out_file | grep -c '20 jobs attempted'";
		my $grep_result = `$grep_cmd`;
		chomp $grep_result;

		if (defined ($grep_result) && $grep_result == 1) {
			$test_result = 1;
		}
	
	}
	
	return $test_result;
}

# ---------------------------------------------------------------------------------------
# Has Rosetta failed many times?
# Check in the DB for whether Rosetta has been started many times (and presumably failed)
# ---------------------------------------------------------------------------------------
sub is_rosetta_always_failing {

	print STDERR "--- check_status.pm: is_rosetta_always_failing\n";

	my $pdb = $_[0];
	my $chain1 = $_[1];
	my $chain2 = $_[2];

	warn "ERROR: PDB code undefined." unless defined $pdb;
	warn "ERROR: Chain 1 undefined."  unless defined $chain1;
	warn "ERROR: Chain 2 undefined."  unless defined $chain2;

	my $dbh = DBI->connect(
		"dbi:SQLite:dbname=workflow_tracker.db",
		"",
		"",
		{ RaiseError => 1 },
	) or die $DBI::errstr;

	# Grab Rosetta attempt count
	my $sql = "SELECT pdb_partner.rosetta_attempts ";
	$sql   .= "FROM pdb, pdb_partner ";
	$sql   .= "WHERE pdb.pdb_id = '$pdb' ";
	$sql   .= "AND pdb_partner.chain_1 = '$chain1' ";
	$sql   .= "AND pdb_partner.chain_2 = '$chain2' ";
	$sql   .= "AND pdb.id = pdb_partner.pdb_id;";
	
	my $sth = $dbh->prepare($sql);
	$sth->execute();

	my $ros_attempts = $sth->fetch();

	my $test_result = 0;
	if (defined (@$ros_attempts) && join ("", @$ros_attempts) >= 3) {
		$test_result = 1;
	}
	
	return $test_result;

	$sth->finish();
	$dbh->disconnect();

}

# ---------------------------------------------------------------------------------------
# Have we downloaded the PDB file? (3)
# Test for whether PDB file exists and its last line is "END"
# ---------------------------------------------------------------------------------------
sub is_pdb_downloaded {
	my $pdb = $_[0];
	
	warn "ERROR: PDB code undefined." unless defined $pdb;
	
	my $last_line;
	my $pdb_file = "pdb_files/" . $pdb . ".pdb";

	my $test_result = 0;
	if (-e $pdb_file) {
		my $tail_cmd = "tail -n1 $pdb_file | grep -c 'END'";
		$test_result = `$tail_cmd`;
	}

	return $test_result;

}

# Have we run find_interface_chains.pl? (4)
# ---------------------------------------------------------------------------------------
# Have we run find_interface_chains.pl? (4)
# Check DB: Is the number of chains > 0 and is the list of chains populated?
# ---------------------------------------------------------------------------------------
sub are_interface_chains_found {
	my $pdb = $_[0];
	
	warn "ERROR: PDB code undefined." unless defined $pdb;
	
	my $dbh = DBI->connect(
		"dbi:SQLite:dbname=workflow_tracker.db",
		"",
		"",
		{ RaiseError => 1 },
	) or die $DBI::errstr;

	# Grab Rosetta attempt count
	my $sql = "SELECT num_chains, chains ";
	$sql   .= "FROM pdb ";
	$sql   .= "WHERE pdb.pdb_id = '$pdb' ";

	my $sth = $dbh->prepare($sql);
	$sth->execute();

	my $chain_info = $sth->fetch();

	my $test_result = 0;
	if (defined (@$chain_info[0])) {
		my $num_chains = @$chain_info[0];
		my $chain_list = @$chain_info[1];

		if (defined ($chain_list) && $num_chains > 0) {
			$test_result = 1;
		}
	}

	return $test_result;

	$sth->finish();
	$dbh->disconnect();

}

# ---------------------------------------------------------------------------------------
# Have we generated the PDB partner file? (5)
# Test for whether the last line of the PDB partner file contains "TER"
# ---------------------------------------------------------------------------------------
sub is_partner_file_generated {
	my $pdb = $_[0];
	my $chain1 = $_[1];
	my $chain2 = $_[2];

	warn "ERROR: PDB code undefined." unless defined $pdb;
	warn "ERROR: Chain 1 undefined."  unless defined $chain1;
	warn "ERROR: Chain 2 undefined."  unless defined $chain2;
	
	my $pdb_part_file = 'pdb_partner_files/' . $pdb;
	$pdb_part_file .= '_' . $chain1 . '_' . $chain2 . '.pdb';

	my $test_result = 0;
	if (-e $pdb_part_file) {
		my $tail_cmd = "tail -n1 $pdb_part_file | grep -c 'TER'";
		$test_result = `$tail_cmd`;
	}

	return $test_result;

}

# ---------------------------------------------------------------------------------------
# Is proc_ala_scan running?
# Search for the job ID (stored in the DB) using qstat
# ---------------------------------------------------------------------------------------
sub is_proc_ala_scan_running {

	my $pdb = $_[0];
	my $chain1 = $_[1];
	my $chain2 = $_[2];

	warn "ERROR: PDB code undefined." unless defined $pdb;
	warn "ERROR: Chain 1 undefined."  unless defined $chain1;
	warn "ERROR: Chain 2 undefined."  unless defined $chain2;

	my $dbh = DBI->connect(
		"dbi:SQLite:dbname=workflow_tracker.db",
		"",
		"",
		{ RaiseError => 1 },
	) or die $DBI::errstr;

	# Grab proc_ala job ID
	my $sql = "SELECT pdb_partner.proc_ala_job_id ";
	$sql   .= "FROM pdb, pdb_partner ";
	$sql   .= "WHERE pdb.pdb_id = '$pdb' ";
	$sql   .= "AND pdb_partner.chain_1 = '$chain1' ";
	$sql   .= "AND pdb_partner.chain_2 = '$chain2' ";
	$sql   .= "AND pdb.id = pdb_partner.pdb_id;";
	
	my $sth = $dbh->prepare($sql);
	$sth->execute();

	my $proc_job_id = $sth->fetch();

	my $job_status;

	# Only do if we have a job ID
	if (@$proc_job_id[0]) {
		my $proc_job_id_str = join ("", @$proc_job_id);

		my $qstat_cmd = "qstat $proc_job_id_str | tail -n1 | awk '{ print \$5 }'";
		$job_status = `$qstat_cmd`;
		chomp $job_status;
	}
	
	my $test_result = 0;
	if (defined ($job_status) && $job_status eq "R") {
		$test_result = 1;
	}
	
	return $test_result;

	$sth->finish();
	$dbh->disconnect();

}

#########################################################################################
### Is proc_ala_scan done? ##############################################################
#########################################################################################

# ---------------------------------------------------------------------------------------
# Has proc_ala_scan failed many times?
# Check in the DB for whether proc_ala_scan has started many times (and presumably failed)
# ---------------------------------------------------------------------------------------
sub is_proc_ala_scan_always_failing {
	my $pdb = $_[0];
	my $chain1 = $_[1];
	my $chain2 = $_[2];

	warn "ERROR: PDB code undefined." unless defined $pdb;
	warn "ERROR: Chain 1 undefined."  unless defined $chain1;
	warn "ERROR: Chain 2 undefined."  unless defined $chain2;

	my $dbh = DBI->connect(
		"dbi:SQLite:dbname=workflow_tracker.db",
		"",
		"",
		{ RaiseError => 1 },
	) or die $DBI::errstr;

	# Grab proc_ala attempt count
	my $sql = "SELECT pdb_partner.proc_ala_attempts ";
	$sql   .= "FROM pdb, pdb_partner ";
	$sql   .= "WHERE pdb.pdb_id = '$pdb' ";
	$sql   .= "AND pdb_partner.chain_1 = '$chain1' ";
	$sql   .= "AND pdb_partner.chain_2 = '$chain2' ";
	$sql   .= "AND pdb.id = pdb_partner.pdb_id;";
	
	my $sth = $dbh->prepare($sql);
	$sth->execute();

	my $proc_ala_attempts = $sth->fetch();

	my $test_result = 0;
	if (defined (@$proc_ala_attempts) && join ("", @$proc_ala_attempts) >= 3) {
		$test_result = 1;
	}
	
	return $test_result;

	$sth->finish();
	$dbh->disconnect();

}

1;